from elasticsearch import Elasticsearch
from elasticsearch_dsl import connections, Index, Document, Text, Integer, Boolean, Keyword, Search, Q, DenseVector
import numpy as np
from openai import OpenAI
import numpy as np
import os
import gpt.gpt_key as gpy_key

os.environ["http_proxy"] = "http://127.0.0.1:7890"
os.environ["https_proxy"] = "http://127.0.0.1:7890"
openai_key = gpy_key.gpt_key

# 连接es数据库
connections = connections.create_connection(hosts=["http://localhost:9200"], timeout=20)
# index = Index('books')
# index.exists() # 判断该索引是否存在
# index.create() # 创建新的索引
# index.get_mapping() # 查看该索引设置的字段
# print(index)
print(connections)


class Questions(Document):
    id = Integer()
    field = Text()
    field_vector = DenseVector(1536)
    original_vector = DenseVector(1536)

    class Index:
        name = "questions"


def get_gpt_embedding(api_key, text):
    """
    实现文本嵌入
    :param api_key: gpt的key
    :param text: 需要计算嵌入的文本
    :return: 向量【1536】的numpy数组
    """
    client = OpenAI(api_key=api_key)
    model = "text-embedding-ada-002"
    emb_req = client.embeddings.create(input=[text], model=model)
    emb = emb_req.data[0].embedding
    return np.array(emb)


test = "你好呀"
test_field = "人与自然"
res = get_gpt_embedding(api_key=openai_key, text=test)
res_field = get_gpt_embedding(api_key=openai_key, text=test_field)
question1 = Questions(id=1, field="人与自然",field_vector=res_field,original_vector=res)
question1.save()

# test = "你好呀"
# res =  get_gpt_embedding(api_key=openai_key,text=test)
# print(res)
# print(res.shape)


# 定义索引和映射
# class Books(Document):
#     title = Text()
#     price = Integer()
#
#     class Index:
#         name = "books"
#
#
# if index.exists():
#     print("删除改索引")
#     index.delete()
#
# book = Books(title="三国演义",price=50)
# book.save()
# book = Books(title="霸道总裁爱上我",price=20)
# book.save()

# print(index.get_mapping())
# # 创建一个查询对象
# search = Search(index='books')
# # 创建一个query
# query = Q('match',title='三国演绎')
# s = search.query(query)
# res = s.execute()
# print(res.hits)
# es_option = Elasticsearch([{'host':'localhost','port':9200}], timeout=3600)
# print(es_option)

# # 查询
# query = {
#   "query": {
#     "match_all": {}
#   }
# }
# result = es_option.search(index="megacorp", body=query)

# 创建索引
# body = {
#   "mappings": {
#     "properties": {
#       "vector": {
#         "type": "dense_vector",
#         "dims": 384,
#         "index": True,
#         "similarity": "l2_norm"
#       },
#       "title": {
#         "type": "text",
#         "fields": {
#           "keyword": {
#             "type": "keyword",
#             "index": False
#           }
#         }
#       },
#       "company": {
#         "type": "keyword",
#         "index": False
#       },
#       "location": {
#         "type": "keyword",
#         "index": False
#       },
#       "salary": {
#         "type": "keyword",
#         "index": False
#       },
#       "job_description": {
#         "type": "keyword",
#         "index": False
#       }
#     }
#   }
# }
#
# # response = es_option.indices.create(index="test",body=body)
# # print(response)
#
# # 查询
# query = {
#   "query": {
#     "match_all": {}
#   }
# }
# result = es_option.search(index="test", body=query)
# print(result)
